TensorFlow 模型优化工具包 — 剪枝 API

自推出 模型优化工具包 （一套新手和高级开发者都可以用来优化机器学习模型的技术）以来，  我们一直忙于践行 路线图，以增加一些新的方法和工具。今天，我们非常高兴地向大家介绍新推出的这款权重剪枝 API。

注：路线图 链接

https://www.tensorflow.org/model_optimization/guide/roadmap

权重剪枝

优化机器学习程序可以采用的形式天差地别。幸运的是，事实证明，神经网络能够适应针对此目标的不同转换。

这类优化工具旨在通过移除神经网络层之间的连接，进而移除参数来减少计算过程中涉及的参数和操作数量。

该权重剪枝 API 基于 Keras 构建，开发者可以非常轻松地将此项技术应用于任何现有的 Keras 训练程序。此 API 将成为适用于该模型优化工具包的新 GitHub 代码库 的一部分，除此之外，我们还将推出许多其他优化技巧（https://github.com/tensorflow/model-optimization）。

import tensorflow_model_optimization as tfmot

model = build_your_model()  

pruning_schedule = tfmot.sparsity.keras.PolynomialDecay(

                        initial_sparsity=0.0, final_sparsity=0.5,

                        begin_step=2000, end_step=4000)

model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule)

...

model_for_pruning.fit(...)

什么是权重剪枝？

权重剪枝指的是消除权重张量中不必要的值。在实践中，我们通过将神经网络参数的值设置为零来移除我们估计存在于神经网络层之间的不必要连接。此操作在训练过程中完成，以使神经网络能够适应这些变化。

权重剪枝为什么有用？

这项工作的直接好处是磁盘压缩，稀疏张量非常适合压缩。因此，通过对剪枝后的 TensorFlow 检查点或转换后的 TensorFlow Lite 模型进行简单的文件压缩，我们可以缩减模型的大小以便存储或传送。例如，在 教程 中，我们演示了如何将一个稀疏度为 90% 的 MNIST 模型从 12MB 压缩到 2MB。

注：教程 链接

https://www.tensorflow.org/model_optimization/guide/pruning/pruning_with_keras

此外，在多个实验中，我们发现权重剪枝与量化兼容，并会产生复合效益。在同一个教程中，我们还展示了如何通过应用 训练后量化 实现进一步压缩，将经过剪枝的模型从 2MB 压缩到仅 0.5MB（https://medium.com/tensorflow/introducing-the-model-optimization-toolkit-for-tensorflow-254aca1ba0a3）。

未来，TensorFlow Lite 会新增对稀疏表征和计算的出色支持，从而将压缩优势扩展到运行时内存，并打破性能改进限制，因为稀疏张量让我们可以跳过涉及零值的不必要计算。

在多个模型上的结果

在实验中，我们验证了此项技术可以成功应用于不同任务的各类模型，适用范围涵盖基于卷积的图像处理神经网络和采用递归神经网络的语音处理模型。下表展示了其中一些实验结果的子集。

在不同模型和任务中获得的稀疏结果。

它的工作原理是什么？

我们基于 Keras 的权重剪枝 API 使用一个非常简单，但又可以广泛应用的 算法，该算法的目的是在训练期间根据模型的数量级，以迭代方式移除连接。从根本上说，开发者可以指定最终目标稀疏度（例如 90%）和执行剪枝的计划（例如从 2000 步开始，每 100 步进行一次剪枝，到 10000 步时停止），以及剪枝结构的可选配置（例如以某种形状应用单个值或值块）。

注：算法 链接

https://arxiv.org/abs/1710.01878

无稀疏度的张量示例（左），1x1 值块中的稀疏度（中），1x2 值块中的稀疏度（右）

随着训练的进行，系统会按计划执行剪枝程序，以消除（即设置为零）拥有最低数量级值（即最接近零的值）的权重，直到达到当前的稀疏度目标。系统每次按计划执行剪枝程序时，都会重新计算当前的稀疏度目标，方法是根据平滑上升函数，从 0% 开始逐渐增加稀疏度，直到在剪枝计划结束时达到最终目标稀疏度。

稀疏度平滑函数示例，剪枝计划从 0 步开始，到 100 步结束，最终目标稀疏度为 90%

和计划一样，开发者可以根据需要调整此平滑函数。例如，在某些情况下，比较方便的做法可能是，安排训练过程在某一步后达到一定的收敛水平时开始优化，或者在完成训练程序中的总步数之前结束剪枝，以便在达到最终目标稀疏度水平时进一步微调系统。如需有关这些配置的更多详细信息，请参阅我们的 教程 和 文档 （https://www.tensorflow.org/model_optimization/guide/pruning）。

注：教程 链接

https://www.tensorflow.org/model_optimization/guide/pruning/pruning_with_keras

在训练过程结束时，根据“经过剪枝”的 Keras 层的最终稀疏度目标，该层的对应张量会包含零。

对张量应用剪枝的动画演示。黑色单元表示存在非零权重的位置。稀疏度随训练的进行而增加

全新文档和 Github 代码库

如前所述，权重剪枝 API 将成为一个新 GitHub 项目和代码库 的一部分，面向能够提高机器学习模型执行和/或表现效率的技术。对于有意踏足这个令人振奋的机器学习领域，或者只是想获得模型优化资源的开发者而言，这是一个很好的开端项目。

注：GitHub 项目和代码库 链接

https://www.tensorflow.org/model_optimization/guide/pruning

鉴于此领域的重要性，我们还将在 tensorflow.org/model_optimization 下创建一个新的子网站，用于提供相关文档和资源。建议您立刻试用此 API，同时欢迎您提供反馈。请关注此博客，跟进日后公告！

注：GitHub 项目和代码库 

http://www.tensorflow.org/model_optimization

致谢

感谢 Raziel Alvarez、Pulkit Bhuwalka、Lawrence Chan、Alan Chiao、Tim Davis、Suyog Gupta、Jian Li、Yunlu Li、Sarah Sirajuddin、Daniel Situnayake、Suharsh Sivakumar。

更多 AI 相关阅读：